import numpy as np
from sklearn.model_selection import train_test_split


# 罪名数量分布
def distribute():
    labels = np.load('data/preprocess/accusation.npy')
    x = labels.sum(axis=1)
    for i in range(int(x.max()) + 1):
        print('罪名数=%d的有%d个' % (i, (x == i).sum()))


def resample(data, min_count, max_times, mode):
    """
    :param data: 原始数据
    :param min_count: 某一类别需填充至的最小数目
    :param max_times: 某一类别可扩充的最大倍数（优先满足）
    :param mode: accusation或relevant_articles
    """
    data_train, data_test = train_test_split(data, test_size=0.05, random_state=1)

    label = []
    label_index = []
    for i, case in enumerate(data_train):
        if len(case['meta'][mode]) == 1:
            label.append(case['meta'][mode][0])
            label_index.append(i)
    label_array = np.array(label)
    label_index_array = np.array(label_index)

    label_set = set(label)
    index_add = []
    add_count = 0
    for i in label_set:
        count = label.count(i)
        index = label_index_array[label_array == i]
        nums = max(int(min_count / count) - 1, 0)
        add_count += (min(max_times, nums) * count)
        if nums > 1:
            index_add += list(index) * min(max_times, nums)
    np.save('data/preprocess/index_add_'+mode+'_%d_%d.npy' % (min_count, max_times), np.array(index_add))

